# -*- coding: utf-8 -*-
"""
Created on Mon Feb  7 13:00:24 2022

@author: user
"""

import asyncio
import json
import numpy as np
import pandas as pd
import traceback
import websockets
import sys
from okex_common import *
from okex_influx import *
from okex_symbol_list import *


global crypto_name, start_time, midp, vol_fresh_count, emar, emar2, emarmr, vol_ema
midp = []
vol_fresh_count = 0
start_time = get_local_timestamp_ms()
emar = {'1min': 0, '5min': 0, '10min': 0, '30min': 0, '60min': 0,
        '1minup': 0, '1mindn': 0, '5minup': 0, '5mindn': 0, '10minup': 0, '10mindn': 0}
emar2 = {'1min': 0, '5min': 0, '10min': 0, '30min': 0, '60min': 0,
         '1minup': 0, '1mindn': 0, '5minup': 0, '5mindn': 0, '10minup': 0, '10mindn': 0}
emarmr = {'1min': 0, '5min': 0, '10min': 0, '30min': 0, '60min': 0,
          '1minup': 0, '1mindn': 0, '5minup': 0, '5mindn': 0, '10minup': 0, '10mindn': 0}
vol_ema = {'1min': 0, '5min': 0, '10min': 0, '30min': 0, '60min': 0,
           '1minup': 0, '1mindn': 0, '5minup': 0, '5mindn': 0, '10minup': 0, '10mindn': 0}
time_trans = 1000000

async def unsubscribe(url, api_key, passphrase, secret_key, channels):
    async with websockets.connect(url) as ws:
        # login
        timestamp = str(get_local_timestamp())
        login_str = login_params(timestamp, api_key, passphrase, secret_key)
        await ws.send(login_str)
        # print(f"send: {login_str}")

        res = await ws.recv()
        print(f"recv: {res}")

        # unsubscribe
        sub_param = {"op": "unsubscribe", "args": channels}
        sub_str = json.dumps(sub_param)
        await ws.send(sub_str)
        print(f"send: {sub_str}")

        res = await ws.recv()
        print(f"recv: {res}")


# unsubscribe channels without login
async def unsubscribe_without_login(url, channels):
    async with websockets.connect(url) as ws:
        # unsubscribe
        sub_param = {"op": "unsubscribe", "args": channels}
        sub_str = json.dumps(sub_param)
        await ws.send(sub_str)
        print(f"send: {sub_str}")

        res = await ws.recv()
        print(f"recv: {res}")
        
        
async def market(url, api_key, passphrase, secret_key, symbol, ordernum):
    global crypto_name, midp, vol_fresh_count, emar, emar2, emarmr, vol_ema
    while True:
        try:
            async with websockets.connect(url) as ws:
                # login
                print(f"{get_timestamp()} [factor] connect success")
                timestamp = str(get_local_timestamp())
                loginStr = login_params(
                    timestamp, api_key, passphrase, secret_key)
                await ws.send(loginStr)
                print(f"{get_timestamp()} [factor] login send: {loginStr}")
                res = await ws.recv()
                print(f"{get_timestamp()} [factor] login recv: {res}")
                
                market_channel = [
                    {"channel": "books50-l2-tbt", "instId": symbol},
                ]
                sub_param = {"op": "subscribe", "args": market_channel}
                sub_str = json.dumps(sub_param)
                await ws.send(sub_str)
                while True:
                    try:
                        res = await asyncio.wait_for(ws.recv(), timeout=20)
                    except (
                        asyncio.TimeoutError,
                        websockets.exceptions.ConnectionClosed,
                    ) as e:
                        try:
                            await ws.send("ping")
                            res = await ws.recv()
                            print(res)
                            continue
                        except Exception as e:
                            print(f"{get_timestamp()} [factor] 连接关闭, 正在重连……")
                            break
                    now = get_local_timestamp_ms()
                    res = eval(res)
                    if "event" in res:
                        # subscribe event message
                        print(
                            "%s [factor] subscrie success, message: %s"
                            % (get_timestamp(), res)
                        )
                    else:
                        if "books50-l2-tbt" == res["arg"]["channel"]:
                            if res["action"] == "snapshot":
                                # 获取首次全量深度数据
                                bids_p, asks_p, instrument_id = partial(res)
                            elif res["action"] == "update":
                                # 获取合并后数据
                                # 获取买的最低价，卖的最高价
                                bids_p = update_bids(res, bids_p)
                                asks_p = update_asks(res, asks_p)
                            else:
                                print(
                                    f"{get_timestamp()} [factor] recv unknown action: {res}"
                                )

                            # 校验checksum
                            server_check_sum = res["data"][0]["checksum"]
                            server_time = int(res["data"][0]["ts"])
                            local_check_num = check(bids_p, asks_p)
                            if local_check_num == server_check_sum:
                                # calculate factors
                                bid = np.array(bids_p)[
                                    :ordernum].T.astype('float')
                                ask = np.array(asks_p)[
                                    :ordernum].T.astype('float')
                                bp = bid[0]
                                bv = bid[1]
                                ap = ask[0]
                                av = ask[1]
                                # write bbo to influx
                                sql = f"{crypto_name}_tbt ap={ap[0]},av={av[0]},bp={bp[0]},bv={bv[0]},server_ts={server_ts} {get_local_timestamp_ms() * time_trans}"
                                # print(sql)
                                write_to_influx(sql)
                                #中间价
                                price = (bp[0] * av[0] + ap[0] *
                                        bv[0]) / (av[0] + bv[0])
                                # 计算短周期波动率
                                midp.append(price)
                                vol_fresh_count = vol_fresh_count + 1
                                if len(midp) > 100:
                                    r_pct = midp[-1] / midp[0] - 1
                                    midp.pop(0)
                                    '''
                                    emar['1min'], emar2['1min'] = ema_std(
                                        r_pct, emar['1min'], emar2['1min'], 6000)
                                    emar['1minup'], emar2['1minup'] = ema_std(
                                        max(r_pct, 0), emar['1minup'], emar2['1minup'], 6000)
                                    emar['1mindn'], emar2['1mindn'] = ema_std(
                                        min(r_pct, 0), emar['1mindn'], emar2['1mindn'], 6000)
                                    emar['5min'], emar2['5min'] = ema_std(
                                        r_pct, emar['5min'], emar2['5min'], 30000)
                                    emar['5minup'], emar2['5minup'] = ema_std(
                                        max(r_pct, 0), emar['5minup'], emar2['5minup'], 30000)
                                    emar['5mindn'], emar2['5mindn'] = ema_std(
                                        min(r_pct, 0), emar['5mindn'], emar2['5mindn'], 30000)
                                    emar['10min'], emar2['10min'] = ema_std(
                                        r_pct, emar['10min'], emar2['10min'], 60000)
                                    emar['10minup'], emar2['10minup'] = ema_std(
                                        max(r_pct, 0), emar['10minup'], emar2['10minup'], 60000)
                                    emar['10mindn'], emar2['10mindn'] = ema_std(
                                        min(r_pct, 0), emar['10mindn'], emar2['10mindn'], 60000)
                                    emar['30min'], emar2['30min'] = ema_std(
                                        r_pct, emar['30min'], emar2['30min'], 180000)
                                    '''
                                    emar['60min'], emar2['60min'] = ema_std(
                                        r_pct, emar['60min'], emar2['60min'], 360000)

                                # 刷新短周期波动率
                                if vol_fresh_count > 3000:
                                    vol_ema = {}
                                    for kname in ['60min']:
                                        vol_ema[kname] = np.sqrt(
                                            emar2[kname] - emar[kname]**2)
                                    vol_fresh_count = 0
                                    print(
                                        f"{get_timestamp()} [factor] vol_ema: {vol_ema['60min']}"
                                    )

                                    # write vol_ema to influx
                                    write_to_influx(
                                        f"{crypto_name}_coef ts={now},vol={vol_ema['60min']} {now * 1000000}"
                                    )
                            else:
                                print(
                                    "%s [factor] 校验结果为: False, 正在重新订阅……"
                                    % (get_timestamp())
                                )
                                # 取消订阅
                                await unsubscribe_without_login(url, market_channel)
                                # 发送订阅
                                async with websockets.connect(url) as ws:
                                    sub_param = {
                                        "op": "subscribe",
                                        "args": market_channel,
                                    }
                                    sub_str = json.dumps(sub_param)
                                    await ws.send(sub_str)
                                    print(f"send: {sub_str}")
                        else:
                            print(
                                f"{get_timestamp()} [factor] recv unknown channel msg: {res}"
                            )
                              
        except Exception as e:
            print(f"{get_timestamp()} [factor] unrecoverable error:")
            print(f"{traceback.format_exc()}")
            print(f"{traceback.format_exc(e)}")
            continue
        
        
if __name__ == "__main__":
    if len(sys.argv) == 1:
        print(f"{get_timestamp()} please specify the crypto name.")
        sys.exit()
    else:
        crypto_name = sys.argv[1]
        okex_symbol, bn_symbol, precision = load_symbol(crypto_name)
        print(f"{get_timestamp()} start {crypto_name} factor download")
    
    global fresh_time
    # 实盘 real trading
    public_url = "wss://ws.okx.com:8443/ws/v5/public"
    ordernum = 1
    freq = 10
    cal_length = 60
    ticknum = 200
    voltick = 100
    #account41
    api_key = "439884fc-57aa-4b61-9b54-c398de0bcc81"
    secret_key = "9FC06FCC4BD7F01ACDACA0FEE659B5C9"
    passphrase = "Unicorn1314"
    print(okex_symbol, bn_symbol)

    loop = asyncio.get_event_loop()

    # # 公共频道 不需要登录（行情, 持仓总量, K线, 标记价格, 深度, 资金费率等）subscribe public channel
    #task_bndata = loop.create_task(bn_market(bn_symbol))
    task_factor = loop.create_task(market(public_url, api_key, passphrase, secret_key, okex_symbol, ordernum))
    task_coefcal = loop.create_task(coef_cal(freq, cal_length, precision, ticknum))
    task_influx = loop.create_task(influx())

    tasks = asyncio.wait(
        [
            task_factor,
            task_coefcal,
            task_influx,
        ]
    )

    loop.run_until_complete(tasks)
    loop.close()
